package utils

import (
	"bytes"
	"crypto/sha256"
	"fmt"
	"math"
)

// Hash computes the SHA256 hash of a byte slice.
func Hash(data []byte) []byte {
	h := sha256.New()
	h.Write(data)
	return h.Sum(nil)
}

// MerkleTree builds a Merkle tree from a list of data stripes.
// Returns the tree as a flat slice and the root hash.
func MerkleTree(stripes [][]byte) ([][]byte, []byte) {
	n := len(stripes)
	if n == 0 {
		return nil, nil
	}

	bottomRowSize := int(math.Pow(2, math.Ceil(math.Log2(float64(n)))))
	tree := make([][]byte, 2*bottomRowSize)

	for i := 0; i < n; i++ {
		tree[bottomRowSize+i] = Hash(stripes[i])
	}

	for i := bottomRowSize - 1; i > 0; i-- {
		left := tree[i*2]
		if left == nil {
			left = []byte{}
		}
		right := tree[i*2+1]
		if right == nil {
			right = []byte{}
		}
		tree[i] = Hash(append(left, right...))
	}

	return tree, tree[1]
}

// GetMerkleBranch returns the authentication path for a leaf at a given index.
func GetMerkleBranch(index int, tree [][]byte) ([][]byte, error) {
	if tree == nil {
		return nil, fmt.Errorf("tree is nil")
	}
	bottomRowSize := len(tree) / 2
	if index < 0 || index >= bottomRowSize {
		return nil, fmt.Errorf("index %d out of bounds for tree with %d leaves", index, bottomRowSize)
	}

	var branch [][]byte
	curr := index + bottomRowSize
	for curr > 1 {
		sibling := curr ^ 1 // Get the sibling's index
		branch = append(branch, tree[sibling])
		curr /= 2
	}
	return branch, nil
}

// MerkleVerify checks if a value is part of the Merkle tree using its branch.
func MerkleVerify(n int, value, rootHash []byte, branch [][]byte, index int) bool {
	if n > 1 && len(branch) != int(math.Ceil(math.Log2(float64(n)))) {
		return false
	}
	if n == 1 && len(branch) != 0 {
		return false
	}

	computedHash := Hash(value)
	for _, siblingHash := range branch {
		if index%2 == 0 { // current is a left child
			computedHash = Hash(append(computedHash, siblingHash...))
		} else { // current is a right child
			computedHash = Hash(append(siblingHash, computedHash...))
		}
		index /= 2
	}

	return bytes.Equal(computedHash, rootHash)
}
